{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook demonstrates the process of adding custom layers to the CoreML model during conversion. We discuss  three examples.\n",
    "\n",
    "For TensorFlow operations (ops for short) that are not translatable to any of the CoreML layers, custom layers can be inserted in the CoreML model (list of CoreML layers can be found [here](https://github.com/apple/coremltools/blob/master/mlmodel/format/NeuralNetwork.proto) or [here](https://apple.github.io/coremltools/coremlspecification/sections/NeuralNetwork.html)). At runtime, CoreML framework will look for the implementation code of the custom layers, which has to be provided by the developer in their app.   \n",
    "Custom layer is a [proto message](https://github.com/apple/coremltools/blob/5b5b8190764ffe78110be6b4d0edbeebe0253a6e/mlmodel/format/NeuralNetwork.proto#L2280), like any other neural network layer in the .mlmodel file (which is in the protobuf format), that can hold the parameters and weights (if any) associated with the TF op.\n",
    "Here is the [documentation](https://developer.apple.com/documentation/coreml/core_ml_api/creating_a_custom_layer) on CoreML custom layers and a nice detailed [blogpost](http://machinethink.net/blog/coreml-custom-layers/). \n",
    "\n",
    "There are two ways in which a custom layer can be added during conversion from TF:\n",
    "\n",
    "1. Specify the argument \"add_custom_layers=True\" during conversion. This will automatically check for unsupported ops and insert a coreml custom layer message in place of that op. The message can be later edited, if required, to add/remove any parameters.            \n",
    "\n",
    "2. Specify the arguments \"add_custom_layers=True\" and \"custom_conversion_functions\" to the converter. The second argument is a dictionary, with keys that are either op types or op names and values are user-defined function handles. The functions receive TensorFlow [op](https://github.com/tensorflow/tensorflow/blob/51ef16057b4625e0a3e2943a9f1bbf856cf098ca/tensorflow/python/framework/ops.py#L3707) object and the CoreML neural network [builder object](https://github.com/apple/coremltools/blob/5b5b8190764ffe78110be6b4d0edbeebe0253a6e/coremltools/models/neural_network.py#L34) and give the user full control on how to handle the TF op and which layers to add to the CoreML graph. When the key is an op type, the function is called whenever op of that type is encountered while traversing the TF graph. Operation names as keys are useful for targeting specific ops. \n",
    "\n",
    "Lets now dive into the examples to make this process clear. \n",
    "\n",
    "First up, setting up some utilities:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "import tensorflow as tf\n",
    "import tensorflow.contrib.slim as slim\n",
    "from tensorflow.python.tools.freeze_graph import freeze_graph\n",
    "import numpy as np\n",
    "import shutil\n",
    "import tempfile\n",
    "import os\n",
    "import tfcoreml\n",
    "import coremltools\n",
    "from coremltools.proto import NeuralNetwork_pb2\n",
    "import netron # we use netron: https://github.com/lutzroeder/Netron for visualization. Comment out this line and all the calls to the \"_visualize\" method, if you do not want to use it. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# A utility function to freeze rhe graph. It will be used later\n",
    "def _simple_run_and_freeze(graph, output_name, frozen_model_file='', feed_dict={}):\n",
    "    \n",
    "    model_dir = tempfile.mkdtemp()\n",
    "    graph_def_file = os.path.join(model_dir, 'tf_graph.pbtxt')\n",
    "    checkpoint_file = os.path.join(model_dir, 'tf_model.ckpt')\n",
    "    \n",
    "    tf.reset_default_graph()\n",
    "    with graph.as_default() as g:\n",
    "      saver = tf.train.Saver()\n",
    "\n",
    "    with tf.Session(graph = graph) as sess:\n",
    "      # initialize\n",
    "      sess.run(tf.global_variables_initializer())\n",
    "      # run the result\n",
    "      fetch = graph.get_operation_by_name(output_name).outputs[0]\n",
    "      tf_result = sess.run(fetch, feed_dict=feed_dict)\n",
    "      # save graph definition somewhere\n",
    "      tf.train.write_graph(sess.graph, model_dir, graph_def_file)\n",
    "      # save the weights\n",
    "      saver.save(sess, checkpoint_file)\n",
    "    \n",
    "    freeze_graph(input_graph=graph_def_file,\n",
    "                 input_saver=\"\",\n",
    "                 input_binary=False,\n",
    "                 input_checkpoint=checkpoint_file,\n",
    "                 output_node_names=output_name,\n",
    "                 restore_op_name=\"save/restore_all\",\n",
    "                 filename_tensor_name=\"save/Const:0\",\n",
    "                 output_graph=frozen_model_file,\n",
    "                 clear_devices=True,\n",
    "                 initializer_nodes=\"\")\n",
    "    \n",
    "    if os.path.exists(model_dir):\n",
    "        shutil.rmtree(model_dir)\n",
    "    \n",
    "    return tf_result\n",
    "\n",
    "# A utility function that takes an MLModel instance and prints info about Neural network layers inside.\n",
    "# It prints short info about all the NN layers and the full description of any custom layer found\n",
    "def _print_coreml_nn_layer_info(spec):\n",
    "    nn_layers = coremltools.models.utils._get_nn_layers(spec)\n",
    "    for i, layer in enumerate(nn_layers):\n",
    "        if layer.WhichOneof('layer') == 'custom':\n",
    "            print( 'layer_id = ', i)\n",
    "            print( layer)\n",
    "        else:\n",
    "            print('{}: layer type: ({}) , inputs: {}, outputs: {}'.\n",
    "              format(i,layer.WhichOneof('layer'), \", \".join([x for x in layer.input]), \", \".join([x for x in layer.output])))\n",
    "\n",
    "# We use \"netron\" for visualization. \n",
    "def _visualize(network_path, port_number):\n",
    "    \n",
    "    def visualize_using_netron(path, port_number):\n",
    "        netron.serve_file(path, browse = True, port=port_number)\n",
    "    \n",
    "    from threading import Thread\n",
    "    import time\n",
    "    \n",
    "    d = Thread(target = visualize_using_netron, args = (network_path, port_number,))\n",
    "    d.setDaemon(True)\n",
    "    d.start()\n",
    "    time.sleep(5)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets define the first TF graph. This one applies a dense layer and normalizes it. It uses the [\"Tile\"](https://www.tensorflow.org/versions/master/api_docs/python/tf/tile) op that CoreML does not support. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define a TF graph: input -> Dense -> unit norm -> output\n",
    "graph = tf.Graph()\n",
    "with graph.as_default() as g:\n",
    "    inputs = tf.placeholder(tf.float32, shape=[None,8], name='input')\n",
    "    with slim.arg_scope([slim.fully_connected],\n",
    "          weights_initializer=tf.truncated_normal_initializer(0.0, 0.2),\n",
    "          weights_regularizer=slim.l2_regularizer(0.0005)):\n",
    "        y = slim.fully_connected(inputs, 10, scope='fc')\n",
    "        y = slim.unit_norm(y,dim=1)\n",
    "\n",
    "output_name = y.op.name\n",
    "X = np.random.rand(1,8)\n",
    "frozen_model_file = 'unit_norm_graph.pb'\n",
    "coreml_model_path = 'unit_norm_graph.mlmodel'\n",
    "out = _simple_run_and_freeze(graph, output_name, frozen_model_file, feed_dict={'input:0' : X})\n",
    "print( 'TF out: ', output_name, out.shape, np.sum(out ** 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize the frozen TF model\n",
    "_visualize(frozen_model_file, np.random.randint(8000, 9000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Try to convert it : this call should raise an error\n",
    "coreml_model = tfcoreml.convert(\n",
    "        tf_model_path=frozen_model_file,\n",
    "        mlmodel_path=coreml_model_path,\n",
    "        input_name_shape_dict={'input:0':[1,8]},\n",
    "        output_feature_names=['UnitNorm/div:0'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# we got an unsupported op error. Try again with custom Flag set to true\n",
    "coreml_model = tfcoreml.convert(\n",
    "        tf_model_path=frozen_model_file,\n",
    "        mlmodel_path=coreml_model_path,\n",
    "        input_name_shape_dict={'input:0':[1,8]},\n",
    "        output_feature_names=['UnitNorm/div:0'],\n",
    "        add_custom_layers=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the \"Tile\" op was made into a custom layer in the CoreML model. This op takes in two inputs, it recasts the first one into the shape given by the second input (by repetition). Here is the [documentation](https://www.tensorflow.org/versions/master/api_docs/python/tf/tile).  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize CoreML model\n",
    "_visualize(coreml_model_path, np.random.randint(8000, 9000))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: As we can see in the visualization, the tensors whose values do not change based on the graph inputs (potentially they depend on the shape of the input, which needs to be fixed during conversion) are converted to \"load constant\" layers in the CoreML graph. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# inspect the CoreML model\n",
    "spec = coreml_model.get_spec()\n",
    "_print_coreml_nn_layer_info(spec)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "\"ClassName\" is an important message: this is the name of the swift/objective-c class that needs to implemented in the Xcode app and will contain the actual code for running the layer.  \n",
    "The \"tile\" op does not have any parameters, so there is no need to edit generated the coreml specification. Lets now convert a TF graph with the op [\"TopKV2\"](https://www.tensorflow.org/api_docs/python/tf/nn/top_k) that has parameters. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define a TF graph: input -> Dense -> softmax -> top_k -> output\n",
    "tf.reset_default_graph()\n",
    "graph = tf.Graph()\n",
    "with graph.as_default() as g:\n",
    "    x = tf.placeholder(tf.float32, shape=[None,8], name=\"input\")\n",
    "    y = tf.layers.dense(inputs=x, units=12, activation=tf.nn.relu)\n",
    "    y = tf.nn.softmax(y, axis=1)\n",
    "    y = tf.nn.top_k(y, k=3, sorted = False, name='output')\n",
    "    \n",
    "output_name = 'output'    \n",
    "X = np.random.rand(1,8)\n",
    "frozen_model_file = 'topk_graph.pb'\n",
    "coreml_model_path = 'topk_graph.mlmodel'\n",
    "out = _simple_run_and_freeze(graph, output_name, frozen_model_file, feed_dict={'input:0' : X})\n",
    "print( 'TF out: ', output_name, out.shape, out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize the frozen TF model\n",
    "_visualize(frozen_model_file, np.random.randint(8000, 9000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Try to convert it : this call should raise an error\n",
    "coreml_model = tfcoreml.convert(\n",
    "        tf_model_path=frozen_model_file,\n",
    "        mlmodel_path=coreml_model_path,\n",
    "        input_name_shape_dict={'input:0':[1,8]},\n",
    "        output_feature_names=['output:0'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# we got an unsupported op error. Try again with custom Flag set to true\n",
    "coreml_model = tfcoreml.convert(\n",
    "        tf_model_path=frozen_model_file,\n",
    "        mlmodel_path=coreml_model_path,\n",
    "        input_name_shape_dict={'input:0':[1,8]},\n",
    "        output_feature_names=['output:0'],\n",
    "        add_custom_layers=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# inspect the CoreML model\n",
    "spec = coreml_model.get_spec()\n",
    "_print_coreml_nn_layer_info(spec)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[top_k](https://www.tensorflow.org/api_docs/python/tf/nn/top_k) operation has two parameters: 'k' and 'sorted'. In the TF graph, the former is received as an additional input by the op and the latter is an op attribute. \n",
    "Let us modify the MLModel spec directly to add these two parameters to this layer. We need to know a little bit about the custom layer's [proto message](https://github.com/apple/coremltools/blob/5b5b8190764ffe78110be6b4d0edbeebe0253a6e/mlmodel/format/NeuralNetwork.proto#L2280) structure to be able to do that. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nn_layers = coremltools.models.utils._get_nn_layers(spec) # get all the layers as a list\n",
    "del nn_layers[3].input[1] # delete the second input: its just the value of k\n",
    "del nn_layers[3].output[1] # delete the second output\n",
    "nn_layers[3].custom.parameters[\"k\"].intValue = 3\n",
    "nn_layers[3].custom.parameters[\"sorted\"].boolValue = False\n",
    "_print_coreml_nn_layer_info(spec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the spec back out\n",
    "coremltools.models.utils.save_spec(spec, coreml_model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize CoreML model\n",
    "_visualize(coreml_model_path, np.random.randint(8000, 9000))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is an alternative way to do the same thing using the \"custom_conversion_functions\" argument: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _convert_topk(**kwargs):\n",
    "    tf_op = kwargs[\"op\"]\n",
    "    coreml_nn_builder = kwargs[\"nn_builder\"]\n",
    "    constant_inputs = kwargs[\"constant_inputs\"]\n",
    "    \n",
    "    params = NeuralNetwork_pb2.CustomLayerParams()\n",
    "    params.className = 'Top_K'\n",
    "    params.description = \"Custom layer that corresponds to the top_k TF op\"\n",
    "    params.parameters[\"sorted\"].boolValue = tf_op.get_attr('sorted')\n",
    "    # get the value of k\n",
    "    k = constant_inputs.get(tf_op.inputs[1].name, 3)\n",
    "    params.parameters[\"k\"].intValue = k\n",
    "    coreml_nn_builder.add_custom(name=tf_op.name,\n",
    "                                input_names=[tf_op.inputs[0].name],\n",
    "                                output_names=[tf_op.outputs[0].name],\n",
    "                                custom_proto_spec=params)\n",
    "\n",
    "coreml_model = tfcoreml.convert(\n",
    "        tf_model_path=frozen_model_file,\n",
    "        mlmodel_path=coreml_model_path,\n",
    "        input_name_shape_dict={'input:0':[1,8]},\n",
    "        output_feature_names=['output:0'],\n",
    "        add_custom_layers=True,\n",
    "        custom_conversion_functions={'TopKV2': _convert_topk})\n",
    "\n",
    "print(\"\\n \\n ML Model layers info: \\n\")\n",
    "# inspect the CoreML model: this should be same as the one we got above\n",
    "spec = coreml_model.get_spec()\n",
    "_print_coreml_nn_layer_info(spec)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "Lets move on to the third and the final example. Now we will encounter an op that is supported but it errors out due to an unsupported coniguration. It is the [Slice](https://www.tensorflow.org/versions/master/api_docs/python/tf/slice) op."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define a TF graph: input -> conv -> slice -> output\n",
    "tf.reset_default_graph()\n",
    "graph = tf.Graph()\n",
    "with graph.as_default() as g:\n",
    "    x = tf.placeholder(tf.float32, shape=[None,10,10,3], name=\"input\")\n",
    "    W = tf.Variable(tf.truncated_normal([1,1,3,5], stddev=0.1))\n",
    "    y = tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME')\n",
    "    y = tf.slice(y, begin=[0,1,1,1], size=[1,2,2,2], name='output')\n",
    "    \n",
    "output_name = 'output'    \n",
    "X = np.random.rand(1,10,10,3)\n",
    "frozen_model_file = 'slice_graph.pb'\n",
    "coreml_model_path = 'slice_graph.mlmodel'\n",
    "out = _simple_run_and_freeze(graph, output_name, frozen_model_file, feed_dict={'input:0' : X})\n",
    "print( 'TF out: ', output_name, out.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize the frozen TF model\n",
    "_visualize(frozen_model_file, np.random.randint(8000, 9000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Try to convert it : this call should raise an error\n",
    "coreml_model = tfcoreml.convert(\n",
    "        tf_model_path=frozen_model_file,\n",
    "        mlmodel_path=coreml_model_path,\n",
    "        input_name_shape_dict={'input:0':[1,10,10,3]},\n",
    "        output_feature_names=['output:0'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This fails, so we provide a custom layer function. Note that this time, the key in the dictionary provided via  \"custom_conversion_functions\" should be same as the op name (\"output\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _convert_slice(**kwargs):\n",
    "    tf_op = kwargs[\"op\"]\n",
    "    coreml_nn_builder = kwargs[\"nn_builder\"]\n",
    "    constant_inputs = kwargs[\"constant_inputs\"]\n",
    "    \n",
    "    params = NeuralNetwork_pb2.CustomLayerParams()\n",
    "    params.className = 'Slice'\n",
    "    params.description = \"Custom layer that corresponds to the slice TF op\"\n",
    "    # get the value of begin\n",
    "    begin = constant_inputs.get(tf_op.inputs[1].name, [0,0,0,0])\n",
    "    size = constant_inputs.get(tf_op.inputs[2].name, [0,0,0,0])\n",
    "    # add begin and size as two repeated weight fields\n",
    "    begin_as_weights = params.weights.add()\n",
    "    begin_as_weights.floatValue.extend(map(float, begin))\n",
    "    size_as_weights = params.weights.add()\n",
    "    size_as_weights.floatValue.extend(map(float, size))\n",
    "    coreml_nn_builder.add_custom(name=tf_op.name,\n",
    "                                input_names=[tf_op.inputs[0].name],\n",
    "                                output_names=[tf_op.outputs[0].name],\n",
    "                                custom_proto_spec=params)\n",
    "\n",
    "coreml_model = tfcoreml.convert(\n",
    "        tf_model_path=frozen_model_file,\n",
    "        mlmodel_path=coreml_model_path,\n",
    "        input_name_shape_dict={'input:0':[1,10,10,3]},\n",
    "        output_feature_names=['output:0'],\n",
    "        add_custom_layers=True,\n",
    "        custom_conversion_functions={'output': _convert_slice}) # dictionary has op name as the key\n",
    "\n",
    "print(\"\\n \\n ML Model layers info: \\n\")\n",
    "# inspect the CoreML model: this should be same as the one we got above\n",
    "spec = coreml_model.get_spec()\n",
    "_print_coreml_nn_layer_info(spec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize CoreML model\n",
    "_visualize(coreml_model_path, np.random.randint(8000, 9000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
